import subprocess
import os
import sys

input_dir = r"D:\train\south-building\images"
output_dir = r"D:\train\south-building\output"

os.makedirs(output_dir, exist_ok=True)

# 设置UTF-8环境变量
env = os.environ.copy()
env["PYTHONUTF8"] = "1"
env["PYTHONIOENCODING"] = "utf-8"

cmd = [
    "ns-process-data", "images",
    "--sfm-tool", "hloc",
    "--feature-type", "superpoint",
    "--matcher-type", "superglue",
    "--data", input_dir,
    "--output-dir", output_dir,
    "--verbose"
]

try:
    # 创建子进程并实时显示输出
    process = subprocess.Popen(
        cmd, 
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,  # 合并标准错误到标准输出
        text=True,
        encoding='utf-8',
        errors='ignore',
        bufsize=1,  # 行缓冲
        universal_newlines=True
    )
    
    # 实时打印输出
    print("开始处理数据...")
    print("-" * 50)
    for line in process.stdout:
        print(line, end='', flush=True)
    
    # 等待进程完成
    return_code = process.wait()
    
    if return_code == 0:
        print("\n" + "-" * 50)
        print("命令执行成功！")
    else:
        print("\n" + "-" * 50)
        print(f"命令执行失败，错误代码：{return_code}")
        
except Exception as e:
    print(f"发生错误：{str(e)}")